{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Model Compilation and Inference using Tensorflow lite runtime\n",
    "\n",
    "In this example notebook, we describe how to take a pre-trained classification model and compile it using ***TF Lite runtime*** to generate deployable artifacts that can be deployed on the target using the ***TF Lite*** interface. \n",
    " \n",
    " - Pre-trained model: `mobilenetv1` model trained on ***ImageNet*** dataset using ***Tensorflow***  \n",
    " \n",
    "In particular, we will show how to\n",
    "- compile the model (during heterogenous model compilation, layers that are supported will be offloaded to the`TI-DSP` and artifacts needed for inference are generated)\n",
    "- enable debug logs\n",
    "- use deny-layer compilation option to isolate possible problematic layers and create additional model subgraphs\n",
    "- use the generated subgraphs artifacts for inference\n",
    "- perform input preprocessing and output postprocessing\n",
    "\n",
    "    \n",
    "## Tensorflow Lite Runtime based work flow\n",
    "\n",
    "The diagram below describes the steps for Tensorflow Lite Runtime based work flow. \n",
    "\n",
    "Note:\n",
    " - The user needs to compile models(sub-graph creation and quantization) on a PC to generate model artifacts.\n",
    " - The generated artifacts can then be used to run inference on the target.\n",
    "\n",
    "<img src=docs/images/tflrt_work_flow_2.png width=\"400\">\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tqdm\n",
    "import cv2\n",
    "import numpy as np\n",
    "import tflite_runtime.interpreter as tflite\n",
    "import shutil \n",
    "from pathlib import Path\n",
    "from scripts.utils import imagenet_class_to_name, download_model\n",
    "from IPython.display import Markdown as md\n",
    "from scripts.utils import loggerWritter\n",
    "from scripts.utils import get_svg_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define utility function to preprocess input images\n",
    "Below, we define a utility function to preprocess images for `mobilenetv1`. This function takes a path as input, loads the image and preprocesses it for generic ***TFLite*** inference. The steps are as follows: \n",
    "\n",
    " 1. load image\n",
    " 2. convert BGR image to RGB\n",
    " 3. scale image so that the short edge is 256 pixels\n",
    " 4. center-crop image to 224x224 pixels\n",
    " 5. apply per-channel pixel scaling and mean subtraction\n",
    "\n",
    "\n",
    "- Note: If you are using a custom model or a model that was trained using a different framework, please remember to define your own utility function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(image_path):\n",
    "    \n",
    "    # read the image using openCV\n",
    "    img = cv2.imread(image_path)\n",
    "    \n",
    "    # convert to RGB\n",
    "    img = img[:,:,::-1]\n",
    "    \n",
    "    # Most of the tflite models are trained using\n",
    "    # 224x224 images. The general rule of thumb\n",
    "    # is to scale the input image while preserving\n",
    "    # the original aspect ratio so that the\n",
    "    # short edge is 256 pixels, and then\n",
    "    # center-crop the scaled image to 224x224\n",
    "    orig_height, orig_width, _ = img.shape\n",
    "    short_edge = min(img.shape[:2])\n",
    "    new_height = (orig_height * 256) // short_edge\n",
    "    new_width = (orig_width * 256) // short_edge\n",
    "    img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)\n",
    "\n",
    "    startx = new_width//2 - (224//2)\n",
    "    starty = new_height//2 - (224//2)\n",
    "    img = img[starty:starty+224,startx:startx+224]\n",
    "    \n",
    "    # apply scaling and mean subtraction.\n",
    "    # if your model is built with an input\n",
    "    # normalization layer, then you might\n",
    "    # need to skip this\n",
    "    img = img.astype('float32')\n",
    "    for mean, scale, ch in zip([128, 128, 128], [0.0078125, 0.0078125, 0.0078125], range(img.shape[2])):\n",
    "            img[:,:,ch] = ((img.astype('float32')[:,:,ch] - mean) * scale)\n",
    "    img = np.expand_dims(img,axis=0)\n",
    "    \n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile the model\n",
    "In this step, we create TFLite runtime with `tidl_model_import_tflite` delegate library to generate artifacts that offload supported portion of the DL model to the TI DSP.\n",
    " - `tidl_delegate` is created with the options below to calibrate the model for 8-bit fixed point inference\n",
    "   \n",
    "    * **tidl_tools_path** - os.getenv('TIDL_TOOLS_PATH'), path to `TIDL` compilation tools \n",
    "    * **artifacts_folder** - folder where all the compilation artifacts needed for inference are stored \n",
    "    * **tensor_bits** - 8 or 16, is the number of bits to be used for quantization \n",
    "    * **accuracy_level** - 1 or 0, the desired accuracy with quantized model\n",
    "    * **advanced_options:calibration_frames** - number of images to be used for calibration\n",
    "    * **advanced_options:calibration_iterations** - number of iterations for advanced calibration\n",
    "    * **debug_level** - 0 -> no debug, 1 -> rt debug prints, >=2 -> increasing levels of debug and trace dump\n",
    "    * **deny_list** force disable offload of a particular operator to TIDL. \n",
    "    \n",
    "- Note: The path to `TIDL` compilation tools and `aarch64` `GCC` compiler is required for model compilation, both of which are accessed by this notebook using predefined environment variables `TIDL_TOOLS_PATH` and `ARM64_GCC_PATH`. The example usage of both the variables is demonstrated in the cell below. \n",
    "- Please refer to TIDL user guide for further advanced options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calib_images = [\n",
    "'sample-images/elephant.bmp',\n",
    "'sample-images/bus.bmp',\n",
    "'sample-images/bicycle.bmp',\n",
    "'sample-images/zebra.bmp',\n",
    "]\n",
    "output_dir = 'custom-artifacts/tflite/mobilenetv1'\n",
    "tflite_model_path = 'models/public/tflite/mobilenet_v1_1.0_224.tflite'\n",
    "download_model(tflite_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compilation knobs  (optional - In case of debugging accuracy)\n",
    "if a model accuracy at 8bits is not good, user's can try compiling same model at 16 bits with accuracy level of 1. This will reduce the performance, but it will give users a good accuracy bar.\n",
    "As a second step, user can try to increase 8 bits accuracy by increasing the number of calibration frames and iterations, in order to get closer to 16 bits + accuracy level of 1 results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "#compilation options - knobs to tweak \n",
    "num_bits =8\n",
    "accuracy =1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Layers debug (optional - In case of debugging)\n",
    "Debug_level 3 gives layer information and warnings/erros which could be useful during debug. User's can see logs from compilation inside a giving path to \"loggerWritter\" helper function.\n",
    "\n",
    "Another technique is to use deny_list to exclude layers from running on TIDL and create additional subgraphs, in order to aisolate issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = Path(\"logs\").mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# stdout and stderr saved to a *.log file.  \n",
    "with loggerWritter(\"logs/custon-model-tfl\"):\n",
    "    \n",
    "# model compilation options\n",
    "    compile_options = {\n",
    "        'tidl_tools_path' : os.environ['TIDL_TOOLS_PATH'],\n",
    "        'artifacts_folder' : output_dir,\n",
    "        'tensor_bits' : num_bits,\n",
    "        'accuracy_level' : accuracy,\n",
    "        'advanced_options:calibration_frames' : len(calib_images),\n",
    "        'advanced_options:calibration_iterations' : 3,\n",
    "        'advanced_options:add_data_convert_ops' : 1,\n",
    "        'debug_level' : 3,\n",
    "        #'deny_list' : 1, #For details of TFLite builtin ops please refer: https://github.com/tensorflow/tensorflow/blob/r2.3/tensorflow/lite/builtin_ops.h\n",
    "    }\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> Please note 'deny_list' is used in above cell as an example and it can be deleted as \"AveragePool2d\" is a supported layers\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the output dir if not preset\n",
    "# clear the directory\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "for root, dirs, files in os.walk(output_dir, topdown=False):\n",
    "    [os.remove(os.path.join(root, f)) for f in files]\n",
    "    [os.rmdir(os.path.join(root, d)) for d in dirs]\n",
    "\n",
    "tidl_delegate = [tflite.load_delegate(os.path.join(os.environ['TIDL_TOOLS_PATH'], 'tidl_model_import_tflite.so'), compile_options)]\n",
    "interpreter = tflite.Interpreter(model_path=tflite_model_path, experimental_delegates=tidl_delegate)\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()\n",
    "\n",
    "for num in tqdm.trange(len(calib_images)):\n",
    "    interpreter.set_tensor(input_details[0]['index'], preprocess(calib_images[num]))\n",
    "    interpreter.invoke()    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subgraphs visualization  (optional - In case of debugging models and subgraps)\n",
    "Running below cell gives links to complete graph and TIDL subgraphs visualizations. This, along with \"deny_list\" feature, explained above, offer tools for potencially checking and isolating issues in NN model layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subgraph_link =get_svg_path(output_dir) \n",
    "for sg in subgraph_link:\n",
    "    hl_text = os.path.join(*Path(sg).parts[4:])\n",
    "    sg_rel = os.path.join('../', sg)\n",
    "    display(md(\"[{}]({})\".format(hl_text,sg_rel))) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use compiled model for inference\n",
    "Then using ***TF Lite*** with the ***`libtidl_tfl_delegate`*** delegate library we run the model and collect benchmark data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "tidl_delegate = [tflite.load_delegate('libtidl_tfl_delegate.so', {'artifacts_folder': output_dir})]\n",
    "interpreter = tflite.Interpreter(model_path=tflite_model_path, experimental_delegates=tidl_delegate)\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()\n",
    "\n",
    "interpreter.set_tensor(input_details[0]['index'], preprocess('sample-images/elephant.bmp'))\n",
    "\n",
    "#Running inference several times to get an stable performance output\n",
    "for i in range(5):\n",
    "    interpreter.invoke()\n",
    "    \n",
    "res = interpreter.get_tensor(output_details[0]['index'])\n",
    "\n",
    "for idx, cls in enumerate(res[0].squeeze()[1:].argsort()[-5:][::-1]):\n",
    "    print('[%d] %s' % (idx, '/'.join(imagenet_class_to_name(cls))))\n",
    "    \n",
    "from scripts.utils import plot_TI_performance_data, plot_TI_DDRBW_data, get_benchmark_output\n",
    "stats = interpreter.get_TI_benchmark_data()\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10,5))\n",
    "plot_TI_performance_data(stats, axis=ax)\n",
    "plt.show()\n",
    "\n",
    "tt, st, rb, wb = get_benchmark_output(stats)\n",
    "print(f'Statistics : \\n Inferences Per Second   : {1000.0/tt :7.2f} fps')\n",
    "print(f' Inference Time Per Image : {tt :7.2f} ms  \\n DDR BW Per Image        : {rb+ wb : 7.2f} MB')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EVM's console logs (optional - in case of inference failure)\n",
    "\n",
    "To copy console logs from EVM to TI EdgeAI Cloud user's workspace, go to: \"Help -> Troubleshooting -> EVM console log\", In TI's EdgeAI Cloud landing page.\n",
    "\n",
    "Alternatevely, from workspace, open/run evm-console-log.ipynb"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
